from .prototypestorage import PrototypeStorage
from utils import LogManager
from utils import RedisConnector
import json
import random


class RedisStorage(PrototypeStorage):
	def __init__(self):
		self.conn = RedisConnector().get_conn()
		self.prefix = "ProxyIPAddress"
		self.log = LogManager()

	def add(self, proxy_data):
		try:
			proxy_data['score'] = proxy_data['time']
			self.conn.hset("{}:{}".format(self.prefix, proxy_data['protocol'].upper()), proxy_data['ip_port'],
						   json.dumps(proxy_data, ensure_ascii = False))
		except Exception as e:
			self.log.log(self.__class__.__name__, "Add Proxy IP {} Error : {}".format(proxy_data, e))

	def delete(self, protocol, ip_port):
		try:
			self.conn.hdel("{}:{}".format(self.prefix, protocol.upper()), ip_port)
		except Exception as e:
			self.log.log(self.__class__.__name__, "Delete Proxy IP {} Error : {}".format(ip_port, e))

	def get(self, protocol):
		try:
			proxy_data = self.conn.hgetall("{}:{}".format(self.prefix, protocol.upper()))
			if proxy_data:
				return json.loads(proxy_data[random.choice(list(proxy_data))])
			return None
		except Exception as e:
			self.log.log(self.__class__.__name__, "Get Proxy IP Error : {}".format(e))

	def set(self, protocol, ip_port, proxy_data):
		try:
			if self.conn.hexists("{}:{}".format(self.prefix, protocol.upper()), ip_port):
				self.conn.hset("{}:{}".format(self.prefix, protocol.upper()), ip_port,
							   json.dumps(proxy_data, ensure_ascii = False))
		except Exception as e:
			self.log.log(self.__class__.__name__, "Set Proxy IP {} Error : {}".format(ip_port, e))

	def get_conn(self):
		return self.conn
